from Data import search_file
import os
import shutil
import fitz  # PyMuPDF


# 主函数，pdf合并
def diagnosis_or_maintenance_pdf(selected_function,filepaths):
    root_path = None
    scg_path = None
    cref_path = None
    # filepaths = search_file.select_folder()
    if filepaths == "-1":
        return "-1"
    else:
        count = 0
        for path in filepaths:
            if "维护命令" in path.split('\\')[-3]:
                if count == 0:
                    root_path = path.rsplit("维护命令", 1)[0]
                    # root_path = f'{path.rsplit("维护命令", 1)[0]}维护命令'
                    scg_path = f'{path.rsplit("维护命令", 1)[0]}维护命令\\00-pdf'
                    # 检查并创建配置指南输出目录
                    if not os.path.exists(scg_path):
                        os.makedirs(scg_path)
                        print(scg_path)
                count = count + 1
            if "诊断命令" in path.split('\\')[-3]:
                if count == 0:
                    # root_path = f'{path.rsplit("诊断命令", 1)[0]}诊断命令'
                    root_path = path.rsplit("诊断命令", 1)[0]
                    cref_path = f'{path.rsplit("诊断命令", 1)[0]}诊断命令\\00-pdf'
                    # 检查并创建配置指南输出目录
                    if not os.path.exists(cref_path):
                        os.makedirs(cref_path)
                        print(cref_path)
                count = count + 1
        # 检查根路径是否存在
        if None == root_path:
            print(f"错误：找不到根路径")
            # print("请检查配置指南与命令参考编号")
            print("程序执行完成!!!")
            return "-1"
        if selected_function == "SCG/CREF合稿":
            directory_mapping(filepaths, scg_path, cref_path)
            convert_to_pdf_format(scg_path, cref_path, root_path)
            print("程序执行完成!!!")
        if selected_function == "大本合稿":
            convert_to_pdf_format(scg_path, cref_path, root_path)
            print("程序执行完成!!!")


# 目录映射
def directory_mapping(filepaths, scg_path, cref_path):
    directory_map = {}
    for filepath in filepaths:
        parts = filepath.split('\\')
        if len(parts) >= 4:
            # 修改目录层级判断
            # if "封面" in parts[-2] or "扉页" in parts[-2]:
            #     first_level = parts[-2]  # 封面和扉页作为一级目录
            #     second_level = None  # 使用文件名作为二级标识
            # else:
            first_level = parts[-3]  # 其他情况按原来的逻辑处理
            second_level = parts[-2]

            # 初始化一级目录
            if first_level not in directory_map:
                directory_map[first_level] = {}
            # 将文件添加到对应的二级目录
            if second_level not in directory_map[first_level]:
                directory_map[first_level][second_level] = []
            directory_map[first_level][second_level].append(filepath)

    # 处理每个目录
    for first_level, second_levels in directory_map.items():
        for second_level, files in second_levels.items():
            if "封面" in second_level and "扉页" in second_level:
                # 将文件copy到00-pdf
                print(f"\033[93m正在处理: {first_level}->{second_level}\033[0m")
                copy_the_cover_and_title_page(files, scg_path, cref_path, first_level)
                # merge_pdf_volumes(files, second_level, scg_path, cref_path)
            if "维护命令" in first_level or "诊断命令" in first_level:
                if "00-pdf" != second_level and "封面" not in second_level and "扉页" not in second_level:
                    print(f"\033[93m正在处理: {first_level}->{second_level}\033[0m")
                    merge_pdf_volumes(files, second_level, scg_path, cref_path, first_level)


# 分册pdf合并
def merge_pdf_volumes(files, second_level, scg_path, cref_path, first_level):
    try:
        merged_pdf = fitz.open()
        page_offset = 0
        all_bookmarks = []

        for file in files:
            try:
                with fitz.open(file) as pdf:
                    if pdf.page_count == 0:
                        continue

                    # 获取文件名作为一级书签
                    file_name = os.path.basename(file).replace('.pdf', '')
                    first_bookmark = [1, file_name, page_offset + 1]
                    all_bookmarks.append(first_bookmark)

                    # 复制文档
                    merged_pdf.insert_pdf(pdf)

                    # 获取原始书签并处理
                    original_toc = pdf.get_toc(simple=False)  # 获取完整书签信息
                    if original_toc:
                        for bookmark in original_toc:
                            level, title, page, dest = bookmark[:4]
                            if dest:
                                # 提取目标页面和位置信息
                                dest_page = dest.get('page', page - 1)
                                dest_top = dest.get('to', fitz.Point(0, 0)).y
                                # 创建新的书签项，包含精确位置信息
                                adjusted_bookmark = [
                                    level + 1,  # 增加层级，因为文件名是一级
                                    title,
                                    dest_page + page_offset + 1,
                                    {
                                        "kind": 1,
                                        "page": dest_page + page_offset,
                                        "to": fitz.Point(0, dest_top),
                                        "zoom": 0
                                    }
                                ]
                                all_bookmarks.append(adjusted_bookmark)

                    page_offset += pdf.page_count

            except Exception as e:
                print(f"处理文件 {file} 时出错: {str(e)}")
                continue

        # 检查并设置书签
        if len(merged_pdf) > 0 and all_bookmarks:
            try:
                merged_pdf.set_toc(all_bookmarks)
                print(f"成功设置 {len(all_bookmarks)} 个书签")
            except Exception as e:
                print(f"设置书签时出错: {str(e)}")

        # 确定输出路径和保存
        if "维护命令" in first_level:
            output_path = os.path.join(scg_path, f"{second_level}.pdf")
        else:
            output_path = os.path.join(cref_path, f"{second_level}.pdf")

        merged_pdf.save(output_path, garbage=4, clean=True, deflate=True)
        merged_pdf.close()
        print(f"已合并PDF文件: {output_path}")

    except Exception as e:
        print(f"合并PDF文件失败: {str(e)}")


# 复制封面扉页
def copy_the_cover_and_title_page(files, scg_path, cref_path, first_level):
    for file in files:
        try:
            # 获取原始文件名
            original_name = os.path.basename(file)

            # 根据文件名和路径类型确定目标路径
            if "诊断命令" in first_level:
                target_dir = cref_path
            elif "维护命令" in first_level:
                target_dir = scg_path
            else:
                print(f"\033[93m警告：无法确定文件类型: {original_name}\033[0m")
                continue

            # print(f"源文件: {file}")
            # print(f"目标目录: {target_dir}")

            # 处理封面和扉页
            if "封面" in original_name or "扉页" in original_name:
                # 直接复制到目标文件，不使用临时文件
                target_file = os.path.join(target_dir, "封面.pdf" if "封面" in original_name else "扉页.pdf")
                try:
                    shutil.copy2(file, target_file)
                    print(f"已复制文件到: {target_file}")
                except Exception as e:
                    print(f"\033[91m直接复制失败: {str(e)}\033[0m")
                    continue

        except Exception as e:
            print(f"\033[91m处理文件失败 {file}: {str(e)}\033[0m")


# 转pdf大本
def convert_to_pdf_format(scg_path, cref_path, path):
    # print(scg_path)
    # print(cref_path)
    try:
        # 处理配置指南和命令参考
        # 使用三目运算符确定文档类型和输入路径
        doc_type = "维护命令" if scg_path is not None else "诊断命令"
        input_path = scg_path if scg_path is not None else cref_path
        print(f"\033[93m正在合并{doc_type}：\033[0m")
        merged_pdf = fitz.open()
        page_offset = 0
        all_bookmarks = []
        current_parent_level = 1

        # 获取所有PDF文件并排序
        pdf_files = []
        cover_file = os.path.join(input_path, "封面.pdf")
        title_page_file = os.path.join(input_path, "扉页.pdf")

        # 添加封面和扉页
        if os.path.exists(cover_file):
            pdf_files.append(cover_file)
        if os.path.exists(title_page_file):
            pdf_files.append(title_page_file)

        # 获取其他PDF文件并按数字排序
        other_files = [f for f in os.listdir(input_path)
                       if f.endswith('.pdf')
                       and f not in ["封面.pdf", "扉页.pdf"]]
        other_files.sort(key=lambda x: int(''.join(filter(str.isdigit, x)) or '0'))
        pdf_files.extend([os.path.join(input_path, f) for f in other_files])

        # 合并所有PDF文件
        for file in pdf_files:
            try:
                with fitz.open(file) as pdf:
                    if pdf.page_count == 0:
                        continue

                    # 获取文件名作为书签
                    file_name = os.path.basename(file).replace('.pdf', '')

                    # 添加文件名作为一级书签
                    first_bookmark = [1, file_name, page_offset + 1]
                    all_bookmarks.append(first_bookmark)
                    current_parent_level = 1

                    # 复制文档
                    merged_pdf.insert_pdf(pdf)

                    # 获取并处理原始书签
                    original_toc = pdf.get_toc(simple=False)
                    if original_toc:
                        for bookmark in original_toc:
                            level, title, page, dest = bookmark[:4]
                            if dest:
                                # 提取目标页面和位置信息
                                dest_page = dest.get('page', page - 1)
                                dest_top = dest.get('to', fitz.Point(0, 0)).y

                                # 调整书签层级
                                adjusted_level = level + current_parent_level

                                # 创建新的书签项
                                adjusted_bookmark = [
                                    adjusted_level,
                                    title,
                                    dest_page + page_offset + 1,
                                    {
                                        "kind": 1,
                                        "page": dest_page + page_offset,
                                        "to": fitz.Point(0, dest_top),
                                        "zoom": 0
                                    }
                                ]
                                all_bookmarks.append(adjusted_bookmark)

                    page_offset += pdf.page_count

            except Exception as e:
                print(f"处理文件 {file} 时出错: {str(e)}")
                continue

        # 设置书签
        if len(merged_pdf) > 0 and all_bookmarks:
            try:
                merged_pdf.set_toc(all_bookmarks)
                print(f"成功设置 {len(all_bookmarks)} 个书签")
            except Exception as e:
                print(f"设置书签时出错: {str(e)}")

        # 保存合并后的文件
        output_path = os.path.join(path, f"{doc_type} (V1.0).pdf")
        merged_pdf.save(output_path, garbage=4, clean=True, deflate=True)
        merged_pdf.close()
        print(f"已生成{doc_type}合集: {output_path}")

    except Exception as e:
        print(f"合并文档失败: {str(e)}")


if __name__ == '__main__':
    diagnosis_or_maintenance_pdf("大本合稿")
